Add graph-safe ops and Inductor heuristics - #106

Summary

Add the graph-safe MUSA compatibility required by CUDA-facing inference paths and bump torchada from 0.1.83 to 0.1.85:

• add MUSA PrivateUse1 overrides for aten::multinomial, aten::log, and aten::log_ on affected torch_musa releases
• reuse CUDA Inductor Triton template heuristics for MUSA matmul templates
• mirror MUSA_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES when both are present
• add focused correctness and CUDA-graph replay coverage for the new operators

torch_musa version boundary

The TorchAda graph-safe compatibility is temporary and is limited to torch_musa versions before 2.11.0.post2:

torch_musa < 2.11.0.post2  -> keep TorchAda compatibility patches
torch_musa >= 2.11.0.post2 -> use torch_musa/ATen native implementations

The existing _TORCH_MUSA_ACCELERATOR_FIX_VERSION = "2.11.0.post2" is reused for this decision. The version is read from the existing torch.musa.__version__ surface; no extra package-version helper or unconditional import torch_musa is added just for version detection.

On torch_musa >= 2.11.0.post2:

• load_cpp_ops() returns before building/loading TorchAda's C++ override extension
• the Inductor heuristic compatibility patch is skipped
• torch_musa/ATen owns log, log_, and multinomial dispatch directly
• no compile-time macro or version check is kept in musa_ops.mu

The currently published vLLM-MUSA image is still based on torch_musa 2.11.0.post1+musa5.2.0, so it remains on the compatibility side of this boundary.

Scope

This PR does not change global CUDA platform semantics:

• does not patch torch.cuda.is_available()
• does not fake torch.cuda.get_device_capability()
• does not globally rewrite torch.load map locations
• does not add a flash-attention import shim
• does not add SGLang-Omni-specific model logic

Testing

Focused validation was run with the PR source overlaid in the MUSA-enabled vLLM-MUSA image:

PYTHONPATH=src python -m pytest -q tests/test_cuda_patching.py \
  -k "VisibleDevicesEnv or InductorTemplateHeuristics or TestAcceleratorModuleWrapper or test_cpp_ops_loaded_on_musa"

36 passed, 202 deselected

Operator correctness and graph replay tests on the same post1 MUSA runtime:

PYTHONPATH=src python -m pytest -q tests/test_log.py tests/test_multinomial.py

8 passed

The version-boundary tests cover the existing post2 accelerator gate, post2 native Inductor registration behavior, local-version suffixes such as +musa5.2.0, and the post2+ no-C++-override behavior. git diff --check and Python syntax compilation also pass.

Downstream image validation

A direct runtime probe in the latest image confirmed:

torch_musa = 2.11.0.post1+musa5.2.0
torch.musa.is_available() = True
torch.musa.device_count() = 1
aten::log PrivateUse1 kernel = present
aten::log_ PrivateUse1 kernel = present
aten::multinomial PrivateUse1 kernel = present

Request-level vLLM compile smoke

A request-level Qwen3-0.6B vLLM compile and GSM8K follow-up was run on MUSA with CompilationMode.VLLM_COMPILE (mode=3) using the compatible image registry.mthreads.com/mcconline/inference/vllm:v0.24.0-ph1-5.2.0-torch2.9.1.post1-20260805.

The service completed compilation, graph capture, a direct completion request, and the official 200-question GSM8K evaluator:

Accuracy: 0.380
Invalid responses: 0.000
All 200 requests returned HTTP 200

The server log contained no runtime exception, backend compiler failure, graph break, or eager fallback during evaluation.

Limitations

• No public torch_musa 2.11.0.post2 wheel/tag is available yet; post2 behavior is covered by version-boundary tests and the code path is fail-safe for unknown versions.
• The 20260819 and 20260824 Qwen3 startup attempts hit a pre-existing vLLM-MUSA is_neox signature mismatch, so request-level evaluation used the compatible 20260805 image.
• The request-level run validates the Python Inductor path and downstream service behavior; it should not be interpreted as C++ override coverage for a post2 wheel.